{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "375gw63MC9Hg"
      },
      "source": [
        "##### Copyright 2020 Google LLC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6gaVU4XhC_rb"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "HxR_1l6_Oyj4"
      },
      "source": [
        "# Adversarial Regularization for Image Classification\n",
        "\n",
        "The core idea of adversarial learning is to train a model with\n",
        "adversarially-perturbed data (called adversarial examples) in addition to the\n",
        "organic training data. The adversarial examples are constructed to intentionally\n",
        "mislead the model into making wrong predictions or classifications. By training\n",
        "with such examples, the model learns to be robust against adversarial\n",
        "perturbation when making predictions.\n",
        "\n",
        "In this tutorial, we illustrate the following procedure of applying adversarial\n",
        "learning to obtain robust models using the Neural Structured Learning framework:\n",
        "\n",
        "1.  Create a neural network as a base model. In this tutorial, the base model is\n",
        "    created with the `tf.keras` functional API; this procedure is compatible\n",
        "    with models created by `tf.keras` sequential and subclassing APIs as well.\n",
        "2.  Wrap the base model with the **`AdversarialRegularization`** wrapper class,\n",
        "    which is provided by the NSL framework, to create a new `tf.keras.Model`\n",
        "    instance. This new model will include the adversarial loss as a\n",
        "    regularization term in its training objective.\n",
        "3.  Convert examples in the training data to feature dictionaries.\n",
        "4.  Train and evaluate the new model.\n",
        "\n",
        "Both the base and the new model will be evaluated against natural and\n",
        "adversarial inputs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Keu58kiDPNDB"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Install the Neural Structured Learning package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "plo5qfYlPKwQ"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet neural-structured-learning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Mi_rYRjCPVuB"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import numpy as np\n",
        "import neural_structured_learning as nsl"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "23fiFpOJP6ab"
      },
      "source": [
        "## Hyperparameters\n",
        "\n",
        "We collect and explain the hyperparameters (in an `HParams` object) for model\n",
        "training and evaluation.\n",
        "\n",
        "Input/Output:\n",
        "\n",
        "*   **`input_shape`**: The shape of the input tensor. Each image is 28-by-28\n",
        "pixels with 1 channel.\n",
        "*   **`num_classes`**: There are a total of 10 classes, corresponding to 10\n",
        "digits [0-9].\n",
        "\n",
        "Model architecture:\n",
        "\n",
        "*   **`conv_filters`**: A list of numbers, each specifying the number of\n",
        "filters in a convolutional layer.\n",
        "*   **`kernel_size`**: The size of 2D convolution window, shared by all\n",
        "convolutional layers.\n",
        "*   **`pool_size`**: Factors to downscale the image in each max-pooling layer.\n",
        "*   **`num_fc_units`**: The number of units (i.e., width) of each\n",
        "fully-connected layer.\n",
        "\n",
        "Training and evaluation:\n",
        "\n",
        "*  **`batch_size`**: Batch size used for training and evaluation.\n",
        "*  **`epochs`**: The number of training epochs.\n",
        "\n",
        "Adversarial learning:\n",
        "\n",
        "*   **`adv_multiplier`**: The weight of adversarial loss in the training\n",
        "    objective, relative to the labeled loss.\n",
        "*   **`adv_step_size`**: The magnitude of adversarial perturbation.\n",
        "*  **`adv_grad_norm`**: The norm to measure the magnitude of adversarial\n",
        "   perturbation.\n",
        "*  **`pgd_iterations`**: The number of iterative steps to take when using PGD.\n",
        "*  **`pgd_epsilon`**: The bounds of the perturbation. PGD will project back to\n",
        "   this epsilon ball when generating the adversary.\n",
        "*  **`clip_value_min`**: Clips the final adversary to be at least as large as\n",
        "   this value. This keeps the perturbed pixel values in a valid domain.\n",
        "*  **`clip_value_max`**: Clips the final adversary to be no larger than this\n",
        "   value. This also keeps the perturbed pixel values in a valid domain.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BVVJErI8P2SK"
      },
      "outputs": [],
      "source": [
        "class HParams(object):\n",
        "  def __init__(self):\n",
        "    self.input_shape = [28, 28, 1]\n",
        "    self.num_classes = 10\n",
        "    self.conv_filters = [32, 64, 64]\n",
        "    self.kernel_size = (3, 3)\n",
        "    self.pool_size = (2, 2)\n",
        "    self.num_fc_units = [64]\n",
        "    self.batch_size = 32\n",
        "    self.epochs = 5\n",
        "    self.adv_multiplier = 0.2\n",
        "    self.adv_step_size = 0.01\n",
        "    self.adv_grad_norm = 'infinity'\n",
        "    self.pgd_iterations = 40\n",
        "    self.pgd_epsilon = 0.2\n",
        "    self.clip_value_min = 0.0\n",
        "    self.clip_value_max = 1.0\n",
        "\n",
        "HPARAMS = HParams()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FcLusOXeRGoF"
      },
      "source": [
        "## MNIST dataset\n",
        "\n",
        "The [MNIST dataset](http://yann.lecun.com/exdb/mnist/) contains grayscale\n",
        "images of handwritten digits (from '0' to '9'). Each image showes one digit at\n",
        "low resolution (28-by-28 pixels). The task involved is to classify images into\n",
        "10 categories, one per digit.\n",
        "\n",
        "Here we load the MNIST dataset from\n",
        "[TensorFlow Datasets](https://www.tensorflow.org/datasets). It handles\n",
        "downloading the data and constructing a `tf.data.Dataset`. The loaded dataset\n",
        "has two subsets:\n",
        "\n",
        "*   `train` with 60,000 examples, and\n",
        "*   `test` with 10,000 examples.\n",
        "\n",
        "Examples in both subsets are stored in feature dictionaries with the following\n",
        "two keys:\n",
        "\n",
        "*   `image`: Array of pixel values, ranging from 0 to 255.\n",
        "*   `label`: Groundtruth label, ranging from 0 to 9."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Mb03tb2_RFtF"
      },
      "outputs": [],
      "source": [
        "datasets = tfds.load('mnist')\n",
        "\n",
        "train_dataset = datasets['train']\n",
        "test_dataset = datasets['test']\n",
        "\n",
        "IMAGE_INPUT_NAME = 'image'\n",
        "LABEL_INPUT_NAME = 'label'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YDXCb187RMO1"
      },
      "source": [
        "To make the model numerically stable, we normalize the pixel values to [0, 1]\n",
        "by mapping the dataset over the `normalize` function. After shuffling training\n",
        "set and batching, we convert the examples to feature tuples `(image, label)`\n",
        "for training the base model. We also provide a function to convert from tuples\n",
        "to dictionaries for later use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "tPWcdkhGRN-F"
      },
      "outputs": [],
      "source": [
        "def normalize(features):\n",
        "  features[IMAGE_INPUT_NAME] = tf.cast(\n",
        "      features[IMAGE_INPUT_NAME], dtype=tf.float32) / 255.0\n",
        "  return features\n",
        "\n",
        "def convert_to_tuples(features):\n",
        "  return features[IMAGE_INPUT_NAME], features[LABEL_INPUT_NAME]\n",
        "\n",
        "def convert_to_dictionaries(image, label):\n",
        "  return {IMAGE_INPUT_NAME: image, LABEL_INPUT_NAME: label}\n",
        "\n",
        "train_dataset = train_dataset.map(normalize).shuffle(10000).batch(HPARAMS.batch_size).map(convert_to_tuples)\n",
        "test_dataset = test_dataset.map(normalize).batch(HPARAMS.batch_size).map(convert_to_tuples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "489axd64RZ5l"
      },
      "source": [
        "## Base model\n",
        "\n",
        "Our base model will be a neural network consisting of 3 convolutional layers\n",
        "follwed by 2 fully-connected layers (as defined in `HPARAMS`). Here we define\n",
        "it using the Keras functional API. Feel free to try other APIs or model\n",
        "architectures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dcipiJgqRolw"
      },
      "outputs": [],
      "source": [
        "def build_base_model(hparams):\n",
        "  \"\"\"Builds a model according to the architecture defined in `hparams`.\"\"\"\n",
        "  inputs = tf.keras.Input(\n",
        "      shape=hparams.input_shape, dtype=tf.float32, name=IMAGE_INPUT_NAME)\n",
        "\n",
        "  x = inputs\n",
        "  for i, num_filters in enumerate(hparams.conv_filters):\n",
        "    x = tf.keras.layers.Conv2D(\n",
        "        num_filters, hparams.kernel_size, activation='relu')(\n",
        "            x)\n",
        "    if i \u003c len(hparams.conv_filters) - 1:\n",
        "      # max pooling between convolutional layers\n",
        "      x = tf.keras.layers.MaxPooling2D(hparams.pool_size)(x)\n",
        "  x = tf.keras.layers.Flatten()(x)\n",
        "  for num_units in hparams.num_fc_units:\n",
        "    x = tf.keras.layers.Dense(num_units, activation='relu')(x)\n",
        "  pred = tf.keras.layers.Dense(hparams.num_classes, activation=None)(x)\n",
        "  # pred = tf.keras.layers.Dense(hparams.num_classes, activation='softmax')(x)\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=pred)\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "tJViTSlgRr_H"
      },
      "outputs": [],
      "source": [
        "base_model = build_base_model(HPARAMS)\n",
        "base_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kUuJ1cNcRw8e"
      },
      "source": [
        "Next we train and evaluate the base model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PuUwF4Y3R1iv"
      },
      "outputs": [],
      "source": [
        "base_model.compile(optimizer='adam',\n",
        "                   loss=tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "                       from_logits=True),\n",
        "                   metrics=['acc'])\n",
        "base_model.fit(train_dataset, epochs=HPARAMS.epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VfEjXHO7R28n"
      },
      "outputs": [],
      "source": [
        "results = base_model.evaluate(test_dataset)\n",
        "named_results = dict(zip(base_model.metrics_names, results))\n",
        "print('\\naccuracy:', named_results['acc'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "HS5WMFrfTpdT"
      },
      "source": [
        "## Adversarial-regularized model\n",
        "\n",
        "Here we show how to incorporate adversarial training into a Keras model with a\n",
        "few lines of code, using the NSL framework. The base model is wrapped to create\n",
        "a new `tf.Keras.Model`, whose training objective includes adversarial\n",
        "regularization.\n",
        "\n",
        "We will train one using the FGSM adversary and one using a stronger PGD\n",
        "adversary.\n",
        "\n",
        "First, we create config objects with relevant hyperparameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DMUkmsk2TpGv"
      },
      "outputs": [],
      "source": [
        "fgsm_adv_config = nsl.configs.make_adv_reg_config(\n",
        "    multiplier=HPARAMS.adv_multiplier,\n",
        "    # With FGSM, we want to take a single step equal to the epsilon ball size,\n",
        "    # to get the largest allowable perturbation.\n",
        "    adv_step_size=HPARAMS.pgd_epsilon,\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    clip_value_min=HPARAMS.clip_value_min,\n",
        "    clip_value_max=HPARAMS.clip_value_max\n",
        ")\n",
        "\n",
        "pgd_adv_config = nsl.configs.make_adv_reg_config(\n",
        "    multiplier=HPARAMS.adv_multiplier,\n",
        "    adv_step_size=HPARAMS.adv_step_size,\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    pgd_iterations=HPARAMS.pgd_iterations,\n",
        "    pgd_epsilon=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=HPARAMS.clip_value_min,\n",
        "    clip_value_max=HPARAMS.clip_value_max\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aRMh8rm5UYCg"
      },
      "source": [
        "Now we can wrap a base model with `AdversarialRegularization`. Here we create \n",
        "new base models (`base_fgsm_model`, `base_pgd_model`), so that the existing one\n",
        "(`base_model`) can be used in later comparison.\n",
        "\n",
        "The returned `adv_model` is a `tf.keras.Model` object, whose training objective\n",
        "includes a regularization term for the adversarial loss. To compute that loss,\n",
        "the model has to have access to the label information (feature `label`), in\n",
        "addition to regular input (feature `image`). For this reason, we convert the\n",
        "examples in the datasets from tuples back to dictionaries. And we tell the\n",
        "model which feature contains the label information via the `label_keys`\n",
        "parameter.\n",
        "\n",
        "We will create two adversarially regularized models: `fgsm_adv_model`\n",
        "(regularized with FGSM) and `pgd_adv_model` (regularized with PGD)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Ce676AY8UhM5"
      },
      "outputs": [],
      "source": [
        "# Create model for FGSM.\n",
        "base_fgsm_model = build_base_model(HPARAMS)\n",
        "# Create FGSM-regularized model.\n",
        "fgsm_adv_model = nsl.keras.AdversarialRegularization(\n",
        "    base_fgsm_model,\n",
        "    label_keys=[LABEL_INPUT_NAME],\n",
        "    adv_config=fgsm_adv_config\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-e7lxy80UvNP"
      },
      "outputs": [],
      "source": [
        "# Create model for PGD.\n",
        "base_pgd_model = build_base_model(HPARAMS)\n",
        "# Create PGD-regularized model.\n",
        "pgd_adv_model = nsl.keras.AdversarialRegularization(\n",
        "    base_pgd_model,\n",
        "    label_keys=[LABEL_INPUT_NAME],\n",
        "    adv_config=pgd_adv_config\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Nci891oNU_C4"
      },
      "outputs": [],
      "source": [
        "# Data for training.\n",
        "train_set_for_adv_model = train_dataset.map(convert_to_dictionaries)\n",
        "test_set_for_adv_model = test_dataset.map(convert_to_dictionaries)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CTbwMW8YVCyu"
      },
      "source": [
        "Next we compile, train, and evaluate the\n",
        "adversarial-regularized model. There might be warnings like\n",
        "\"Output missing from loss dictionary,\" which is fine because\n",
        "the `adv_model` doesn't rely on the base implementation to\n",
        "calculate the total loss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-SmhMHhkVEfm"
      },
      "outputs": [],
      "source": [
        "fgsm_adv_model.compile(optimizer='adam',\n",
        "                       loss=tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "                        from_logits=True),\n",
        "                       metrics=['acc'])\n",
        "fgsm_adv_model.fit(train_set_for_adv_model, epochs=HPARAMS.epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NUrv3hwmVV0E"
      },
      "outputs": [],
      "source": [
        "results = fgsm_adv_model.evaluate(test_set_for_adv_model)\n",
        "named_results = dict(zip(fgsm_adv_model.metrics_names, results))\n",
        "print('\\naccuracy:', named_results['sparse_categorical_accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nyXwibwdVcYY"
      },
      "outputs": [],
      "source": [
        "pgd_adv_model.compile(optimizer='adam',\n",
        "                       loss=tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "                        from_logits=True),\n",
        "                       metrics=['acc'])\n",
        "pgd_adv_model.fit(train_set_for_adv_model, epochs=HPARAMS.epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sTfLmGyc_VTJ"
      },
      "outputs": [],
      "source": [
        "results = pgd_adv_model.evaluate(test_set_for_adv_model)\n",
        "named_results = dict(zip(pgd_adv_model.metrics_names, results))\n",
        "print('\\naccuracy:', named_results['sparse_categorical_accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oAwpl_IkWTsV"
      },
      "source": [
        "Both adversarially regularized models perform well on the test set."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "iLShGGJ6WXll"
      },
      "source": [
        "## Robustness under Adversarial Perturbations\n",
        "\n",
        "Now we compare the base model and the adversarial-regularized model for\n",
        "robustness under adversarial perturbation.\n",
        "\n",
        "We will show how the base model is vulnerable to attacks from both FGSM and PGD,\n",
        "the FGSM-regularized model can resist FGSM attacks but is vulnerable to PGD, and\n",
        "the PGD-regularized model is able to resist both forms of attack.\n",
        "\n",
        "We use `gen_adv_neighbor` to generate adversaries for our models."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DxoVVC3evZ4C"
      },
      "source": [
        "### Attacking the Base Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UyMn0wLAEWuQ"
      },
      "outputs": [],
      "source": [
        "# Set up the neighbor config for FGSM.\n",
        "fgsm_nbr_config = nsl.configs.AdvNeighborConfig(\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    adv_step_size=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=0.0,\n",
        "    clip_value_max=1.0,\n",
        ")\n",
        "\n",
        "# The labeled loss function provides the loss for each sample we pass in. This\n",
        "# will be used to calculate the gradient.\n",
        "labeled_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True,\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "jpM-5en5vY59"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate adversarial images using FGSM on the base model.\n",
        "perturbed_images, labels, predictions = [], [], []\n",
        "\n",
        "# We want to record the accuracy.\n",
        "metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "for batch in test_set_for_adv_model:\n",
        "  # Record the loss calculation to get the gradient.\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(batch)\n",
        "    losses = labeled_loss_fn(batch[LABEL_INPUT_NAME],\n",
        "                             base_model(batch[IMAGE_INPUT_NAME]))\n",
        "    \n",
        "  # Generate the adversarial example.\n",
        "  fgsm_images, _ = nsl.lib.adversarial_neighbor.gen_adv_neighbor(\n",
        "      batch[IMAGE_INPUT_NAME],\n",
        "      losses,\n",
        "      fgsm_nbr_config,\n",
        "      gradient_tape=tape\n",
        "  )\n",
        "\n",
        "  # Update our accuracy metric.\n",
        "  y_true = batch['label']\n",
        "  y_pred = base_model(fgsm_images)\n",
        "  metric(y_true, y_pred)\n",
        "\n",
        "  # Store images for later use.\n",
        "  perturbed_images.append(fgsm_images)\n",
        "  labels.append(y_true.numpy())\n",
        "  predictions.append(tf.argmax(y_pred, axis=-1).numpy())\n",
        "\n",
        "print('%s model accuracy: %f' % ('base', metric.result().numpy()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vQL2qWI8-h8G"
      },
      "source": [
        "Let's examine what some of these images look like."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "MpMs1syT_p0o"
      },
      "outputs": [],
      "source": [
        "def examine_images(perturbed_images, labels, predictions, model_key):\n",
        "  batch_index = 0\n",
        "\n",
        "  batch_image = perturbed_images[batch_index]\n",
        "  batch_label = labels[batch_index]\n",
        "  batch_pred = predictions[batch_index]\n",
        "\n",
        "  batch_size = HPARAMS.batch_size\n",
        "  n_col = 4\n",
        "  n_row = (batch_size + n_col - 1) / n_col\n",
        "\n",
        "  print('accuracy in batch %d:' % batch_index)\n",
        "  print('%s model: %d / %d' %\n",
        "        (model_key, np.sum(batch_label == batch_pred), batch_size))\n",
        "\n",
        "  plt.figure(figsize=(15, 15))\n",
        "  for i, (image, y) in enumerate(zip(batch_image, batch_label)):\n",
        "    y_base = batch_pred[i]\n",
        "    plt.subplot(n_row, n_col, i+1)\n",
        "    plt.title('true: %d, %s: %d' % (y, model_key, y_base), color='r'\n",
        "      if y != y_base else 'k')\n",
        "    plt.imshow(tf.keras.preprocessing.image.array_to_img(image), cmap='gray')\n",
        "    plt.axis('off')\n",
        "\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kr_lkBtM966D"
      },
      "outputs": [],
      "source": [
        "examine_images(perturbed_images, labels, predictions, 'base')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y8V2ssdF-md7"
      },
      "source": [
        "Our perturbation budget of 0.2 is quite large, but even so, the perturbed\n",
        "numbers are clearly recognizable to the human eye. On the other hand, our\n",
        "network is fooled into misclassifying several examples.\n",
        "\n",
        "As we can see, the FGSM attack is already highly effective, and quick to\n",
        "execute, heavily reducing the model accuracy. We will see below, that the PGD\n",
        "attack is even more effective, even with the same perturbation budget."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "qMbG-KbjGFlB"
      },
      "outputs": [],
      "source": [
        "# Set up the neighbor config for PGD.\n",
        "pgd_nbr_config = nsl.configs.AdvNeighborConfig(\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    adv_step_size=HPARAMS.adv_step_size,\n",
        "    pgd_iterations=HPARAMS.pgd_iterations,\n",
        "    pgd_epsilon=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=HPARAMS.clip_value_min,\n",
        "    clip_value_max=HPARAMS.clip_value_max,\n",
        ")\n",
        "\n",
        "# pgd_model_fn generates a prediction from which we calculate the loss, and the\n",
        "# gradient for a given interation.\n",
        "pgd_model_fn = base_model\n",
        "\n",
        "# We need to pass in the loss function for repeated calculation of the gradient.\n",
        "pgd_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, \n",
        ")\n",
        "labeled_loss_fn = pgd_loss_fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "U0Og8WuKyUGt"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate adversarial images using PGD on the base model.\n",
        "perturbed_images, labels, predictions = [], [], []\n",
        "\n",
        "# Record the accuracy.\n",
        "metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "for batch in test_set_for_adv_model:\n",
        "  # Gradient tape to calculate the loss on the first iteration.\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(batch)\n",
        "    losses = labeled_loss_fn(batch[LABEL_INPUT_NAME],\n",
        "                             base_model(batch[IMAGE_INPUT_NAME]))\n",
        "    \n",
        "  # Generate the adversarial examples.\n",
        "  pgd_images, _ = nsl.lib.adversarial_neighbor.gen_adv_neighbor(\n",
        "      batch[IMAGE_INPUT_NAME],\n",
        "      losses,\n",
        "      pgd_nbr_config,\n",
        "      gradient_tape=tape,\n",
        "      pgd_model_fn=pgd_model_fn,\n",
        "      pgd_loss_fn=pgd_loss_fn,\n",
        "      pgd_labels=batch[LABEL_INPUT_NAME],\n",
        "  )\n",
        "\n",
        "  # Update our accuracy metric.\n",
        "  y_true = batch['label']\n",
        "  y_pred = base_model(pgd_images)\n",
        "  metric(y_true, y_pred)\n",
        "\n",
        "  # Store images for visualization.\n",
        "  perturbed_images.append(pgd_images)\n",
        "  labels.append(y_true.numpy())\n",
        "  predictions.append(tf.argmax(y_pred, axis=-1).numpy())\n",
        "\n",
        "print('%s model accuracy: %f' % ('base', metric.result().numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "F7KJaU4f_LLx"
      },
      "outputs": [],
      "source": [
        "examine_images(perturbed_images, labels, predictions, 'base')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "69naYIftz0Nm"
      },
      "source": [
        "The PGD attack is much stronger, but it also takes longer to run."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Rfj9P6fvz97d"
      },
      "source": [
        "### Attacking the FGSM Regularized Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4qgnKmh9Gy1q"
      },
      "outputs": [],
      "source": [
        "# Set up the neighbor config.\n",
        "fgsm_nbr_config = nsl.configs.AdvNeighborConfig(\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    adv_step_size=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=0.0,\n",
        "    clip_value_max=1.0,\n",
        ")\n",
        "\n",
        "# The labeled loss function provides the loss for each sample we pass in. This\n",
        "# will be used to calculate the gradient.\n",
        "labeled_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XAmIbcPJ0C6d"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate adversarial images using FGSM on the regularized model.\n",
        "perturbed_images, labels, predictions = [], [], []\n",
        "\n",
        "# Record the accuracy.\n",
        "metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "for batch in test_set_for_adv_model:\n",
        "  # Record the loss calculation to get its gradients.\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(batch)\n",
        "    # We attack the adversarially regularized model.\n",
        "    losses = labeled_loss_fn(batch[LABEL_INPUT_NAME],\n",
        "                             fgsm_adv_model.base_model(batch[IMAGE_INPUT_NAME]))\n",
        "    \n",
        "  # Generate the adversarial examples.\n",
        "  fgsm_images, _ = nsl.lib.adversarial_neighbor.gen_adv_neighbor(\n",
        "      batch[IMAGE_INPUT_NAME],\n",
        "      losses,\n",
        "      fgsm_nbr_config,\n",
        "      gradient_tape=tape\n",
        "  )\n",
        "\n",
        "  # Update our accuracy metric.\n",
        "  y_true = batch['label']\n",
        "  y_pred = fgsm_adv_model.base_model(fgsm_images)\n",
        "  metric(y_true, y_pred)\n",
        "\n",
        "  # Store images for visualization.\n",
        "  perturbed_images.append(fgsm_images)\n",
        "  labels.append(y_true.numpy())\n",
        "  predictions.append(tf.argmax(y_pred, axis=-1).numpy())\n",
        "\n",
        "print('%s model accuracy: %f' % ('base', metric.result().numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PO8w0Gmv_XbH"
      },
      "outputs": [],
      "source": [
        "examine_images(perturbed_images, labels, predictions, 'fgsm_reg')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SCB48teL0O48"
      },
      "source": [
        "As we can see, the FGSM-regularized model performs much better than the base\n",
        "model on images perturbed by FGSM. How does it do against PGD?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "fPBmdviOHmlT"
      },
      "outputs": [],
      "source": [
        "# Set up the neighbor config for PGD.\n",
        "pgd_nbr_config = nsl.configs.AdvNeighborConfig(\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    adv_step_size=HPARAMS.adv_step_size,\n",
        "    pgd_iterations=HPARAMS.pgd_iterations,\n",
        "    pgd_epsilon=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=HPARAMS.clip_value_min,\n",
        "    clip_value_max=HPARAMS.clip_value_max,\n",
        ")\n",
        "\n",
        "# pgd_model_fn generates a prediction from which we calculate the loss, and the\n",
        "# gradient for a given interation.\n",
        "pgd_model_fn = fgsm_adv_model.base_model\n",
        "\n",
        "# We need to pass in the loss function for repeated calculation of the gradient.\n",
        "pgd_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, \n",
        ")\n",
        "labeled_loss_fn = pgd_loss_fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "km7Tlg_b0d5Y"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate adversarial images using PGD on the FGSM-regularized model.\n",
        "perturbed_images, labels, predictions = [], [], []\n",
        "\n",
        "metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "for batch in test_set_for_adv_model:\n",
        "  # Gradient tape to calculate the loss on the first iteration.\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(batch)\n",
        "    losses = labeled_loss_fn(batch[LABEL_INPUT_NAME],\n",
        "                             fgsm_adv_model.base_model(batch[IMAGE_INPUT_NAME]))\n",
        "    \n",
        "  # Generate the adversarial examples.\n",
        "  pgd_images, _ = nsl.lib.adversarial_neighbor.gen_adv_neighbor(\n",
        "      batch[IMAGE_INPUT_NAME],\n",
        "      losses,\n",
        "      pgd_nbr_config,\n",
        "      gradient_tape=tape,\n",
        "      pgd_model_fn=pgd_model_fn,\n",
        "      pgd_loss_fn=pgd_loss_fn,\n",
        "      pgd_labels=batch[LABEL_INPUT_NAME],\n",
        "  )\n",
        "  \n",
        "  # Update our accuracy metric.\n",
        "  y_true = batch['label']\n",
        "  y_pred = fgsm_adv_model.base_model(pgd_images)\n",
        "  metric(y_true, y_pred)\n",
        "\n",
        "  # Store images for visualization.\n",
        "  perturbed_images.append(pgd_images)\n",
        "  labels.append(y_true.numpy())\n",
        "  predictions.append(tf.argmax(y_pred, axis=-1).numpy())\n",
        "\n",
        "print('%s model accuracy: %f' % ('base', metric.result().numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "K5i9BB1fAIQB"
      },
      "outputs": [],
      "source": [
        "examine_images(perturbed_images, labels, predictions, 'fgsm_reg')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XQZTIOmn2Ylt"
      },
      "source": [
        "While the FGSM regularized model was robust to attacks via FGSM, as we can see\n",
        "it is still vulnerable to attacks from PGD, which is a stronger attack mechanism\n",
        "than FGSM.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "S6pb7D4E2ixj"
      },
      "source": [
        "### Attacking the PGD Regularized Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5MFkg3lsIYNt"
      },
      "outputs": [],
      "source": [
        "# Set up the neighbor config.\n",
        "fgsm_nbr_config = nsl.configs.AdvNeighborConfig(\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    adv_step_size=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=0.0,\n",
        "    clip_value_max=1.0,\n",
        ")\n",
        "\n",
        "# The labeled loss function provides the loss for each sample we pass in. This\n",
        "# will be used to calculate the gradient.\n",
        "labeled_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9eee2wAM2mBL"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate adversarial images using FGSM on the regularized model.\n",
        "perturbed_images, labels, predictions = [], [], []\n",
        "\n",
        "# Record the accuracy.\n",
        "metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "for batch in test_set_for_adv_model:\n",
        "  # Record the loss calculation to get its gradients.\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(batch)\n",
        "    # We attack the adversarially regularized model.\n",
        "    losses = labeled_loss_fn(batch[LABEL_INPUT_NAME],\n",
        "                             pgd_adv_model.base_model(batch[IMAGE_INPUT_NAME]))\n",
        "\n",
        "  # Generate the adversarial examples.\n",
        "  fgsm_images, _ = nsl.lib.adversarial_neighbor.gen_adv_neighbor(\n",
        "      batch[IMAGE_INPUT_NAME],\n",
        "      losses,\n",
        "      fgsm_nbr_config,\n",
        "      gradient_tape=tape\n",
        "  )\n",
        "\n",
        "  # Update our accuracy metric.\n",
        "  y_true = batch['label']\n",
        "  y_pred = pgd_adv_model.base_model(fgsm_images)\n",
        "  metric(y_true, y_pred)\n",
        "\n",
        "  # Store images for visualization.\n",
        "  perturbed_images.append(fgsm_images)\n",
        "  labels.append(y_true.numpy())\n",
        "  predictions.append(tf.argmax(y_pred, axis=-1).numpy())\n",
        "\n",
        "print('%s model accuracy: %f' % ('base', metric.result().numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wLIrI20WAbyw"
      },
      "outputs": [],
      "source": [
        "examine_images(perturbed_images, labels, predictions, 'pgd_reg')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "SAqEans_I10c"
      },
      "outputs": [],
      "source": [
        "# Set up the neighbor config for PGD.\n",
        "pgd_nbr_config = nsl.configs.AdvNeighborConfig(\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm,\n",
        "    adv_step_size=HPARAMS.adv_step_size,\n",
        "    pgd_iterations=HPARAMS.pgd_iterations,\n",
        "    pgd_epsilon=HPARAMS.pgd_epsilon,\n",
        "    clip_value_min=HPARAMS.clip_value_min,\n",
        "    clip_value_max=HPARAMS.clip_value_max,\n",
        ")\n",
        "\n",
        "# pgd_model_fn generates a prediction from which we calculate the loss, and the\n",
        "# gradient for a given interation.\n",
        "pgd_model_fn = pgd_adv_model.base_model\n",
        "\n",
        "# We need to pass in the loss function for repeated calculation of the gradient.\n",
        "pgd_loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, \n",
        ")\n",
        "labeled_loss_fn = pgd_loss_fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Xzb9qW1K2z_8"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "# Generate adversarial images using PGD on the PGD-regularized model.\n",
        "perturbed_images, labels, predictions = [], [], []\n",
        "\n",
        "metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "for batch in test_set_for_adv_model:\n",
        "  # Gradient tape to calculate the loss on the first iteration.\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(batch)\n",
        "    losses = labeled_loss_fn(batch[LABEL_INPUT_NAME],\n",
        "                             pgd_adv_model.base_model(batch[IMAGE_INPUT_NAME]))\n",
        "  \n",
        "  # Generate the adversarial examples.\n",
        "  pgd_images, _ = nsl.lib.adversarial_neighbor.gen_adv_neighbor(\n",
        "      batch[IMAGE_INPUT_NAME],\n",
        "      losses,\n",
        "      pgd_nbr_config,\n",
        "      gradient_tape=tape,\n",
        "      pgd_model_fn=pgd_model_fn,\n",
        "      pgd_loss_fn=pgd_loss_fn,\n",
        "      pgd_labels=batch[LABEL_INPUT_NAME],\n",
        "  )\n",
        "  \n",
        "  # Update our accuracy metric.\n",
        "  y_true = batch['label']\n",
        "  y_pred = pgd_adv_model.base_model(pgd_images)\n",
        "  metric(y_true, y_pred)\n",
        "\n",
        "  # Store images for visualization.\n",
        "  perturbed_images.append(pgd_images)\n",
        "  labels.append(y_true.numpy())\n",
        "  predictions.append(tf.argmax(y_pred, axis=-1).numpy())\n",
        "\n",
        "print('%s model accuracy: %f' % ('base', metric.result().numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RB2rUM4LAmMM"
      },
      "outputs": [],
      "source": [
        "examine_images(perturbed_images, labels, predictions, 'pgd_reg')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5P29jDqi3cgz"
      },
      "source": [
        "The PGD-regularized model is strong against both attack types."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "luOZc_1B_oq7"
      },
      "source": [
        "# Conclusion\n",
        "\n",
        "In this colab, we've explored two gradient-based attack methods, FGSM, and its\n",
        "stronger variant PGD. We have seen how neural networks not trained to defend\n",
        "against these attacks are vulnerable to these attacks, and also how to utilize\n",
        "adversarial regularization in the Neural Structured Learning framework to\n",
        "improve robustness."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Adversarial Regularization for MNIST Image Classification",
      "private_outputs": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
